# build_middle.py
# usage: 根据questions和results 生成给与二线的middle文件
# python ./src/tools/build_middle.py --question ./data/test.jsonl --result ./result/results.jsonl --middle ./data/middle.{序号范围标识}.txt

import argparse
import json

def main():
    parser = argparse.ArgumentParser(description="Merge question and result files into middle format.")
    parser.add_argument('--question', required=True, help='Path to the question JSONL file.')
    parser.add_argument('--result', required=True, help='Path to the result JSONL file.')
    parser.add_argument('--middle', required=True, help='Path to the output middle JSONL file.')

    args = parser.parse_args()

    # 读取 question 文件
    questions = {}
    with open(args.question, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            data = json.loads(line)
            questions[data['id']] = {
                'level': data['level'],
                'question': data['question']
            }

    # 读取 result 文件
    results = {}
    with open(args.result, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            data = json.loads(line)
            results[data['id']] = data['sql']

    # 合并并写入输出文件
    with open(args.middle, 'w', encoding='utf-8') as out_f:
        # 获取所有 id 并排序（可选，保证输出有序）
        all_ids = sorted(set(questions.keys()) & set(results.keys()))
        
        for id_ in all_ids:
            item = {
                "level": questions[id_]['level'],
                "id": id_,
                "flag": 0,
                "question": questions[id_]['question'],
                "sql": results[id_]
            }
            out_f.write(json.dumps(item, ensure_ascii=False) + '\n')

    print(f"Successfully wrote merged data to {args.middle}")

if __name__ == '__main__':
    main()